{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CqWhSlfDqmns"
   },
   "source": [
    "<a id='prerequisits'></a>\n",
    "\n",
    "# Prerequisits\n",
    "\n",
    "This section installs `gradslam` (if not already installed), imports the necessary packages for the tutorial, and downloads 'lr kt1' (the first trajectory) of [ICL-NUIM dataset](https://www.doc.ic.ac.uk/~ahanda/VaFRIC/iclnuim.html) and structures it as below:\n",
    "```\n",
    "ICL\n",
    "    living_room_traj1_frei_png\n",
    "        depth/    rgb/    associations.txt    livingRoom1.gt.freiburg    livingRoom1n.gt.sim\n",
    "```\n",
    "\n",
    "\n",
    "We set the ICL path variable: `icl_path='ICL/'`. The ICL data is loaded into the following variables: <br>\n",
    "\n",
    "* `colors`: of shape (batch_size, sequence_length, height, width, 3) <br>\n",
    "* `depths`: of shape (batch_size, sequence_length, height, width, 1) <br>\n",
    "* `intrinsics`: of shape (batch_size, 1, 4, 4) <br>\n",
    "* `poses`: of shape (batch_size, sequence_length, 4, 4) <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9opFP8sDqmnt",
    "outputId": "cd17ee21-79da-4b24-ac52-f90b133f80d1"
   },
   "outputs": [],
   "source": [
    "# install gradslam (if not installed)\n",
    "try:\n",
    "    import gradslam as gs\n",
    "except ImportError:\n",
    "    print(\"Installing gradslam...\")\n",
    "    !pip install 'git+https://github.com/gradslam/gradslam.git' -q\n",
    "    print('Installed')\n",
    "\n",
    "# import necessary packages\n",
    "import gradslam as gs\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "from gradslam import Pointclouds, RGBDImages\n",
    "from gradslam.datasets import ICL\n",
    "from gradslam.slam import PointFusion\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# download 'lr kt1' of ICL dataset\n",
    "if not os.path.isdir('ICL'):\n",
    "    os.mkdir('ICL')\n",
    "if not os.path.isdir('ICL/living_room_traj1_frei_png'):\n",
    "    print('Downloading ICL/living_room_traj1_frei_png dataset...')\n",
    "    os.mkdir('ICL/living_room_traj1_frei_png')\n",
    "    !wget http://www.doc.ic.ac.uk/~ahanda/living_room_traj1_frei_png.tar.gz -P ICL/living_room_traj1_frei_png/ -q\n",
    "    !tar -xzf ICL/living_room_traj1_frei_png/living_room_traj1_frei_png.tar.gz -C ICL/living_room_traj1_frei_png/\n",
    "    !rm ICL/living_room_traj1_frei_png/living_room_traj1_frei_png.tar.gz\n",
    "    !wget https://www.doc.ic.ac.uk/~ahanda/VaFRIC/livingRoom1n.gt.sim -P ICL/living_room_traj1_frei_png/ -q\n",
    "    print('Downloaded.')\n",
    "icl_path = 'ICL/'\n",
    "\n",
    "# load dataset\n",
    "dataset = ICL(icl_path, seqlen=8, height=240, width=320)\n",
    "loader = DataLoader(dataset=dataset, batch_size=2)\n",
    "colors, depths, intrinsics, poses, *_ = next(iter(loader))\n",
    "\n",
    "# create rgbdimages object\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "rgbdimages.plotly(0).update_layout(autosize=False, height=600, width=400).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z7KZsY2lqLe1"
   },
   "source": [
    "# Instantiation\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "The `RGBDImages` structures aims to contain batched frame tensors to more easily pass on to SLAM algorithms. It also supports easy computation of (both local and global) vertex maps and normal maps.\n",
    "\n",
    "An `RGBDImages` object can be initialized from rgb images, depth images, instrinsics and (optionally) poses. `RGBDImages` supports both a channels first and a channels last representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VekYeIKkqbsF",
    "outputId": "f1e7a9d5-4d34-462c-f5dd-8ef606c6b3ee"
   },
   "outputs": [],
   "source": [
    "print(f\"colors shape: {colors.shape}\")  # torch.Size([2, 8, 240, 320, 3])\n",
    "print(f\"depths shape: {depths.shape}\")  # torch.Size([2, 8, 240, 320, 1])\n",
    "print(f\"intrinsics shape: {intrinsics.shape}\")  # torch.Size([2, 1, 4, 4])\n",
    "print(f\"poses shape: {poses.shape}\")  # torch.Size([2, 8, 4, 4])\n",
    "print('---')\n",
    "\n",
    "# instantiation without poses\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics)\n",
    "print(rgbdimages.shape)  # (2, 8, 240, 320)\n",
    "print(rgbdimages.poses)  # None\n",
    "print('---')\n",
    "\n",
    "# instantiation with poses\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "print(rgbdimages.shape)  # (2, 8, 240, 320)\n",
    "print('---')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4aAK9HbIqmnv"
   },
   "source": [
    "# Indexing and slicing\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "Basic indexing and slicing of `RGBDImages` over the first (batch) dimension and the second (sequence length) dimension is supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xlqE42usqmnv",
    "outputId": "11057604-571b-4e66-b077-7b94dcda5d61"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "\n",
    "# indexing\n",
    "rgbdimages0 = rgbdimages[0, 0]\n",
    "print(rgbdimages0.shape)  # (1, 1, 240, 320)\n",
    "print('---')\n",
    "\n",
    "# slicing\n",
    "rgbdimages1 = rgbdimages[:2, :5]\n",
    "print(rgbdimages1.shape)  # (2, 5, 240, 320)\n",
    "print('---')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hPIkUUNqqmnx"
   },
   "source": [
    "# Vertex maps and normal maps\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "This section demonstrates accessing vertex maps and normal maps from `RGBDImages`. Vertex maps are computed when accessing the `RGBDImages.vertex_maps` property, and are cached afterwards for additional access without further computation (and similarly with normal maps).\n",
    "\n",
    "`RGBDImages` has both a local vertex map property (`RGBDImages.vertex_map`) which computes vertex positions with respect to each frame, as well as global vertex map (`RGBDImages.global_vertex_map`) which considers the poses of the `RGBDImages` object to compute the global vertex positions. A similar story is true for `RGBDimages.normal_map` and `RGBDImages.global_normal_map`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 263
    },
    "id": "warZULSpqmny",
    "outputId": "3e8cb902-9644-4b45-d32e-f12146ca07dd"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "\n",
    "# compute vertex maps and normal maps\n",
    "print(rgbdimages.vertex_map.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print(rgbdimages.normal_map.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print(rgbdimages.global_vertex_map.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print(rgbdimages.global_normal_map.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print('---')\n",
    "\n",
    "# visualize\n",
    "fig, ax = plt.subplots(1, 2)\n",
    "ax[0].title.set_text('local normal map')\n",
    "ax[0].imshow((rgbdimages.normal_map[-1, -1].numpy() * 255).astype(np.uint8))\n",
    "ax[1].title.set_text('global normal map')\n",
    "ax[1].imshow((rgbdimages.global_normal_map[-1, -1].numpy() * 255).astype(np.uint8))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qfpvR58Sqmn0"
   },
   "source": [
    "# Transfer between GPU/CPU\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "`RGBDImages` support easy transfer between CPU and GPU. This operation transfers all tensors in the `RGBDImages` objects between CPU/GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EbxiaFwSuNvv",
    "outputId": "ff599c26-8614-4280-99da-a8b311ec9e97"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    # transfer to GPU\n",
    "    rgbdimages = rgbdimages.to(\"cuda\")\n",
    "    rgbdimages = rgbdimages.cuda()  # equivalent to rgbdimages.to(\"cuda\")\n",
    "    print(rgbdimages.rgb_image.device)  # \"cuda:0\"\n",
    "    print('---')\n",
    "\n",
    "# transfer to CPU\n",
    "rgbdimages = rgbdimages.to(\"cpu\")\n",
    "rgbdimages = rgbdimages.cpu()  # equivalent to rgbdimages.to(\"cpu\")\n",
    "print(rgbdimages.rgb_image.device)  # \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yrG8sexkuRBK"
   },
   "source": [
    "# Detach and clone tensors\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "`RGBDImages.detach` returns a new `RGBDImages` object such that all internal tensors of the new object do not require grad. `RGBDImages.clone()` returns a new `RGBDImages` object such that all the internal tensors are cloned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q9vkSwS1uWNc",
    "outputId": "7167155f-37b8-4775-c0b4-d9095bec4b3c"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors.requires_grad_(True),\n",
    "                        depths.requires_grad_(True),\n",
    "                        intrinsics.requires_grad_(True),\n",
    "                        poses.requires_grad_(True))\n",
    "\n",
    "# clone\n",
    "rgbdimages1 = rgbdimages.clone()\n",
    "print(torch.allclose(rgbdimages1.rgb_image, rgbdimages.rgb_image))  # True\n",
    "print(rgbdimages1.rgb_image is rgbdimages.rgb_image)  # False\n",
    "print('---')\n",
    "\n",
    "# detach\n",
    "rgbdimages2 = rgbdimages.detach()\n",
    "print(rgbdimages.rgb_image.requires_grad)  # True\n",
    "print(rgbdimages2.rgb_image.requires_grad)  # False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wu3Bbxl-mRt3"
   },
   "source": [
    "# Channels first and channels last representation\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "`RGBDImages` supports both a channels first and a channels last representation. These representations can be transformed to one another with `to_channels_first()` and `to_channels_last()` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3APDidk5mjZu",
    "outputId": "d3bbffe6-361b-4f25-c1e8-13cff7d0343f"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "print(rgbdimages.rgb_image.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print('---')\n",
    "\n",
    "# convert to channels first representation\n",
    "rgbdimages1 = rgbdimages.to_channels_first()\n",
    "print(rgbdimages1.rgb_image.shape)  # torch.Size([2, 8, 3, 240, 320])\n",
    "print('---')\n",
    "\n",
    "# convert to channels last representation\n",
    "rgbdimages2 = rgbdimages1.to_channels_last()\n",
    "print(rgbdimages2.rgb_image.shape)  # torch.Size([2, 8, 240, 320, 3])\n",
    "print('---')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lhHwfuA2mj4j"
   },
   "source": [
    "# Visualization\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "For easy and quick visualization of `RGBDImages`, one can use the `.plotly(batch_index)` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "7iUau5mvmqFV",
    "outputId": "697e4b85-9970-43e6-95c8-3c718721a7d0"
   },
   "outputs": [],
   "source": [
    "# initalize RGBDImages\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "\n",
    "# visualize\n",
    "rgbdimages.plotly(0).show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "rgbdimages_tutorial.ipynb",
   "provenance": []
  },
  "jupytext": {
   "text_representation": {
    "extension": ".py",
    "format_name": "light",
    "format_version": "1.5",
    "jupytext_version": "1.6.0"
   }
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
